#!/usr/bin/python
# -*- coding: UTF-8 -*-

import json
import re
import warnings
import nltk
import pandas as pd

warnings.filterwarnings("ignore", category=DeprecationWarning)

pattern_escape = re.compile(r'&\w{2,3}')
pattern_at = re.compile(r'@\w+')
pattern_marker = re.compile(r"[^a-zA-Z#]")
pattern_url = re.compile('(https?|ftp|file)://[-A-Za-z0-9+&@#/%?=~_|!:,.;]+[-A-Za-z0-9+&@#/%=~_|]')
lemma = nltk.wordnet.WordNetLemmatizer()
stemmer = nltk.stem.porter.PorterStemmer()

def lower(s):
    s = s.lower()  # 返回小写
    return s


def remove_url(s):
    s = pattern_url.sub('', s)  ## 去除url
    return s


def remove_escape(s):
    s = pattern_escape.sub('', s)
    return s


def remove_at(s):
    s = pattern_at.sub('', s)  # 去除@xx
    return s


def remove_marker(s):
    s = pattern_marker.sub(' ', s)  # 仅保留英文单词
    return s


def remove_short(s):
    ss = s.split()
    results = []
    for x in ss:
        if len(x) > 2:
            results.append(x)
    new_line = ' '.join(results)
    return new_line


def do_lemmatize(s):  # 词性还原
    ss = s.split()
    results = []
    for x in ss:
        results.append(lemma.lemmatize(x))
    new_line = ' '.join(results)
    return new_line


def do_stem(s):  # 词干提取
    ss = s.split()
    results = []
    for x in ss:
        results.append(stemmer.stem(x))
    new_line = ' '.join(results)
    return new_line


def clean2():
    covid = 'covid19_tweets.csv'
    df = pd.read_csv(covid)
    df = df[['user_name', 'user_location', 'date', 'text']]
    df = df.dropna(subset=['text', 'user_location'])

    df['text'] = df['text'].apply(lower)
    df['text'] = df['text'].apply(remove_url)
    df['text'] = df['text'].apply(remove_escape)
    df['text'] = df['text'].apply(remove_at)
    df['text'] = df['text'].apply(remove_marker)
    df['text'] = df['text'].apply(remove_short)
    df['text'] = df['text'].apply(do_lemmatize)
    #     df['text'] = df['text'].do_stem(clean_txt)
    df['text'] = df['text'].apply(remove_short)
    df = df.drop(df[df['text'] == ''].index)
    df.to_csv('clean.csv')
    return df


if __name__ == '__main__':
    df = clean2()